import { TABLE_WORKFLOW_SNAPSHOT, normalizePerPage, WorkflowsStorage } from '@mastra/core/storage';
import type { StorageListWorkflowRunsInput, StorageWorkflowRun, WorkflowRun, WorkflowRuns } from '@mastra/core/storage';
import type { StepResult, WorkflowRunState, WorkflowRunStatus } from '@mastra/core/workflows';

import type { StoreOperationsConvex } from '../operations';

type RawWorkflowRun = Omit<StorageWorkflowRun, 'createdAt' | 'updatedAt' | 'snapshot'> & {
  createdAt: string;
  updatedAt: string;
  snapshot: WorkflowRunState | string;
};

export class WorkflowsConvex extends WorkflowsStorage {
  constructor(private readonly operations: StoreOperationsConvex) {
    super();
  }

  async updateWorkflowResults({
    workflowName,
    runId,
    stepId,
    result,
    requestContext,
  }: {
    workflowName: string;
    runId: string;
    stepId: string;
    result: StepResult<any, any, any, any>;
    requestContext: Record<string, any>;
  }): Promise<Record<string, StepResult<any, any, any, any>>> {
    const run = await this.getRun(workflowName, runId);
    if (!run) return {};

    const snapshot = this.ensureSnapshot(run);
    snapshot.context = snapshot.context || {};
    snapshot.context[stepId] = result;
    snapshot.requestContext = { ...(snapshot.requestContext || {}), ...requestContext };

    await this.persistWorkflowSnapshot({
      workflowName,
      runId,
      resourceId: run.resourceId,
      snapshot,
    });

    return JSON.parse(JSON.stringify(snapshot.context));
  }

  async updateWorkflowState({
    workflowName,
    runId,
    opts,
  }: {
    workflowName: string;
    runId: string;
    opts: {
      status: WorkflowRunStatus;
      result?: StepResult<any, any, any, any>;
      error?: string;
      suspendedPaths?: Record<string, number[]>;
      waitingPaths?: Record<string, number[]>;
    };
  }): Promise<WorkflowRunState | undefined> {
    const run = await this.getRun(workflowName, runId);
    if (!run) return undefined;

    const snapshot = this.ensureSnapshot(run);
    const updated: WorkflowRunState = { ...snapshot, ...opts };

    await this.persistWorkflowSnapshot({
      workflowName,
      runId,
      resourceId: run.resourceId,
      snapshot: updated,
    });

    return updated;
  }

  async persistWorkflowSnapshot({
    workflowName,
    runId,
    resourceId,
    snapshot,
  }: {
    workflowName: string;
    runId: string;
    resourceId?: string;
    snapshot: WorkflowRunState;
  }): Promise<void> {
    const now = new Date();
    // Check if a record already exists to preserve createdAt
    const existing = await this.operations.load<{ createdAt?: string } | null>({
      tableName: TABLE_WORKFLOW_SNAPSHOT,
      keys: { workflow_name: workflowName, run_id: runId },
    });

    await this.operations.insert({
      tableName: TABLE_WORKFLOW_SNAPSHOT,
      record: {
        workflow_name: workflowName,
        run_id: runId,
        resourceId,
        snapshot,
        createdAt: existing?.createdAt ?? now.toISOString(),
        updatedAt: now.toISOString(),
      },
    });
  }

  async loadWorkflowSnapshot({
    workflowName,
    runId,
  }: {
    workflowName: string;
    runId: string;
  }): Promise<WorkflowRunState | null> {
    const row = await this.operations.load<{ snapshot: WorkflowRunState | string } | null>({
      tableName: TABLE_WORKFLOW_SNAPSHOT,
      keys: { workflow_name: workflowName, run_id: runId },
    });

    if (!row) return null;
    return typeof row.snapshot === 'string' ? JSON.parse(row.snapshot) : JSON.parse(JSON.stringify(row.snapshot));
  }

  async listWorkflowRuns(args: StorageListWorkflowRunsInput = {}): Promise<WorkflowRuns> {
    const { workflowName, fromDate, toDate, perPage, page, resourceId, status } = args;

    let rows = await this.operations.queryTable<RawWorkflowRun>(TABLE_WORKFLOW_SNAPSHOT, undefined);

    if (workflowName) rows = rows.filter(run => run.workflow_name === workflowName);
    if (resourceId) rows = rows.filter(run => run.resourceId === resourceId);
    if (fromDate) rows = rows.filter(run => new Date(run.createdAt).getTime() >= fromDate.getTime());
    if (toDate) rows = rows.filter(run => new Date(run.createdAt).getTime() <= toDate.getTime());
    if (status) {
      rows = rows.filter(run => {
        const snapshot = this.ensureSnapshot(run);
        return snapshot.status === status;
      });
    }

    const total = rows.length;
    rows.sort((a, b) => new Date(b.createdAt).getTime() - new Date(a.createdAt).getTime());

    if (perPage !== undefined && page !== undefined) {
      const normalized = normalizePerPage(perPage, Number.MAX_SAFE_INTEGER);
      const offset = page * normalized;
      rows = rows.slice(offset, offset + normalized);
    }

    const runs: WorkflowRun[] = rows.map(run => ({
      workflowName: run.workflow_name,
      runId: run.run_id,
      snapshot: this.ensureSnapshot(run),
      createdAt: new Date(run.createdAt),
      updatedAt: new Date(run.updatedAt),
      resourceId: run.resourceId,
    }));

    return { runs, total };
  }

  async getWorkflowRunById({
    runId,
    workflowName,
  }: {
    runId: string;
    workflowName?: string;
  }): Promise<WorkflowRun | null> {
    const runs = await this.operations.queryTable<RawWorkflowRun>(TABLE_WORKFLOW_SNAPSHOT, undefined);
    const match = runs.find(run => run.run_id === runId && (!workflowName || run.workflow_name === workflowName));
    if (!match) return null;

    return {
      workflowName: match.workflow_name,
      runId: match.run_id,
      snapshot: this.ensureSnapshot(match),
      createdAt: new Date(match.createdAt),
      updatedAt: new Date(match.updatedAt),
      resourceId: match.resourceId,
    };
  }

  private async getRun(workflowName: string, runId: string): Promise<RawWorkflowRun | null> {
    const runs = await this.operations.queryTable<RawWorkflowRun>(TABLE_WORKFLOW_SNAPSHOT, [
      { field: 'workflow_name', value: workflowName },
    ]);
    return runs.find(run => run.run_id === runId) ?? null;
  }

  private ensureSnapshot(run: { snapshot: WorkflowRunState | string }): WorkflowRunState {
    if (!run.snapshot) {
      return {
        context: {},
        activePaths: [],
        activeStepsPath: {},
        timestamp: Date.now(),
        suspendedPaths: {},
        resumeLabels: {},
        serializedStepGraph: [],
        value: {},
        waitingPaths: {},
        status: 'pending',
        runId: '',
      };
    }

    if (typeof run.snapshot === 'string') {
      return JSON.parse(run.snapshot);
    }

    return JSON.parse(JSON.stringify(run.snapshot));
  }
}
